const WebSocket = require('ws');
const http = require('http');
const fs = require('fs');
const path = require('path');
const { OpenAI } = require("openai");
const dotenv = require('dotenv');
dotenv.config();

const config = {
    clientOptions: {},
    chatModel: 'gpt-3.5-turbo',
};
if (process.env['OPENAI_API_KEY']) config.clientOptions.apiKey = process.env['OPENAI_API_KEY'];
if (process.env['OPENAI_BASE_URL']) config.clientOptions.baseURL = process.env['OPENAI_BASE_URL'];
if (process.env['OPENAI_CHAT_MODEL']) config.chatModel = process.env['OPENAI_CHAT_MODEL'];

const openai = new OpenAI(config.clientOptions);
const chatModel = config.chatModel;

const server = http.createServer((req, res) => {

    const parsedUrl = new URL(req.url, `http://${req.headers.host}`);
    const pathname = parsedUrl.pathname;
    const filepath = path.join(__dirname, '/client', pathname.substring(-1) === '/' ? pathname + 'index.html' : pathname);

    fs.access(filepath, fs.constants.F_OK, (err) => {
        if (err) {
            res.writeHead(404, { 'Content-Type': 'text/plain' });
            res.end('404 Not Found\n');
            return;
        }

        fs.readFile(filepath, (err, data) => {
            if (err) {
                res.writeHead(500, { 'Content-Type': 'text/plain' });
                res.end('500 Internal Server Error\n');
                return;
            }

            const mimeType = getMimeType(filepath);
            res.writeHead(200, { 'Content-Type': mimeType });
            res.end(data);
        });
    });
});

const wss = new WebSocket.Server({ server });
wss.on('connection', (ws) => {
    console.log('Client connected');

    ws.on('message', async (data) => {
        const message = data.toString();
        console.log(`Received message: ${message}`);

        const stream = await openai.chat.completions.create({
            model: chatModel,
            messages: [{ role: 'user', content: message }],
            stream: true,
        });
        console.log("stream start");
        for await (const chunk of stream) {
            const text = chunk.choices[0].delta?.content || '';
            // console.log("chunk text", text);
            if (text) {
                ws.send(text);
            }
        }
        console.log("stream end");
        ws.send('\0');
    });

    ws.on('close', () => {
        console.log('Client disconnected');
    });
});

server.listen(8080, () => {
    console.log('Server is listening on port 8080');
});

// util function to get mime type
function getMimeType(filepath) {
    const extname = path.extname(filepath);
    switch (extname) {
      case '.html':
        return 'text/html';
      case '.css':
        return 'text/css';
      case '.js':
        return 'application/javascript';
      case '.json':
        return 'application/json';
      case '.png':
        return 'image/png';
      case '.jpg':
        return 'image/jpeg';
      // 你可以根据需要添加更多的MIME类型
      default:
        return 'application/octet-stream'; // 二进制流，用于未知文件类型
    }
}